package chainx

import (
	"context"
	"fmt"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/prompts"
	"strings"
)

type Reasoning struct {
	*chains.LLMChain
	OutThinkKey string
}

func NewReasoning(llm llms.Model, prompt prompts.FormatPrompter, opts ...chains.ChainCallOption) *Reasoning {
	return &Reasoning{
		LLMChain:    chains.NewLLMChain(llm, prompt, opts...),
		OutThinkKey: "think",
	}
}

func (c Reasoning) Call(ctx context.Context, values map[string]any, options ...chains.ChainCallOption) (map[string]any, error) {
	res, err := c.LLMChain.Call(ctx, values, options...)
	if err != nil {
		return res, err
	}
	out, ok := res[c.OutputKey].(string)
	if !ok {
		return res, fmt.Errorf("out is not a string out : %v", out)
	}

	arr := strings.Split(out, "</think>")
	if len(arr) < 2 {
		return res, fmt.Errorf("out not inexistence </think> : %v", out)
	}
	thinkArr := strings.Split(arr[0], "<think>")
	if len(thinkArr) < 2 {
		return res, fmt.Errorf("out not inexistence <think> : %v", out)
	}

	res[c.OutputKey] = strings.Trim(arr[1], "\n")
	res[c.OutThinkKey] = strings.Trim(thinkArr[1], "\n")
	return res, nil
}
